"""
Reinforcement learning maze example.

Red rectangle:          explorer.
Black rectangles:       hells       [reward = -1].
Yellow bin circle:      paradise    [reward = +1].
All other states:       ground      [reward = 0].

This script is the main part which controls the update method of this example.
The RL is in RL_brain.py.

View more on my tutorial page: https://morvanzhou.github.io/tutorials/
"""

from maze_env import Maze
from RL_brain import QLearningTable


def update():
    for episode in range(1000): # 学习 100 回合
        if episode < 100: # 初始回合随机探路, 回合起始位置也随机
            rand = True
        else:
            rand = False
        observation = env.reset(rand) # 初始化 state 的观测值
        setp_n = 0

        while True:
            env.render() # 更新可视化环境

            # RL 大脑根据 state 的观测值挑选 action
            action = RL.choose_action(observation, rand)

            # 探索者在环境中实施这个 action, 并得到环境返回的下一个 state 观测值, reward 和 done (是否是掉下地狱或者升上天堂)
            observation_, reward, done = env.step(action)

            # RL 从这个序列 (state, action, reward, state_) 中学习
            RL.learn(observation, action, done, reward, observation_)

            # 将下一个 state 的值传到下一次循环
            observation = observation_

            setp_n += 1
            if done: # 如果掉下黑洞，或者达到目标, 这回合就结束了
                if reward < 0: # 打印出达到目标的回合状态
                    print('episode:',episode,'setp_n:',setp_n,'reward:',reward)
                break

    print('game over')
    env.destroy() # 结束游戏并关闭窗口

if __name__ == "__main__":
    # 定义环境 env 和 RL 方式
    env = Maze()
    RL = QLearningTable(actions=list(range(env.n_actions))) #, env.MAZE_W, env.MAZE_H)

    # 开始可视化环境 env
    env.after(100, update)
    env.mainloop()
